/**
 *  @file SSLSocket.cpp
 */
#include "SSLSocket.h"
#include "SSLException.h"
#include "SSLSocketException.h"
#include "../../common/UtilFuncT.h"
#include "../../common/ObjPtrT.h"
#include "../../lang/Primitive.h"
#include "../../security/cert/X509Certificate.h"
#include "../__internal/SocketToggleBlockingState.h"

#include "openssl/bio.h" 
#include "openssl/err.h" 
#include "openssl/ssl.h" 
//#include "openssl/x509v3.h"

using namespace cppflib::lang;
using namespace cppflib::security::cert;

namespace cppflib
{

namespace net
{

namespace ssl
{


PRIVATE SSL *__GetSSL(ptrany_t sslObj)
{
   return static_cast<SSL*>(sslObj);
}

SSLSocket::SSLSocket(void)
{
   this->sockAttr->SetAttr(Socket::SSL_CHK_VERIFYRESULT, true);
   this->sockAttr->SetAttr(Socket::SSL_SENDRECV_WITHSSL, true);
}

SSLSocket::SSLSocket(Socket::sockattr_t verifyFlag)
{
   this->sockAttr->SetAttr(verifyFlag, true);
   this->sockAttr->SetAttr(Socket::SSL_SENDRECV_WITHSSL, true);
}

SSLSocket::~SSLSocket(void)
{
   Close();
}

/////////////////// Methods present in Socket but not supported by SSLSocket /////////////////

bool SSLSocket::Create(const CFString &socketAddress, u16_t socketPort, Socket::SocketType socketType)
{
   throw SSLSocketException(_S("Don't call this version of 'Create' from a SSLSocket"));
   return false;
}

bool SSLSocket::Accept(Socket &connectedSocket)
{
   throw SSLSocketException(_S("Don't call this version of 'Accept' from a SSLSocket"));
   return false;
}

Socket * SSLSocket::Accept()
{
   throw SSLSocketException(_S("Don't call this version of 'Accept' from a SSLSocket"));
   return NULL;
}

bool SSLSocket::Listen(int connectionBacklog)
{
   throw SSLSocketException(_S("Don't call 'Listen' from a SSLSocket"));
   return false;
}

int SSLSocket::ReceiveFrom(u8_t * pBuf, int bufLen, int nFlags, CFString &peerAddress, u16_t &peerPort)
{
   throw SSLSocketException(_S("Don't call 'ReceiveFrom' from a SSLSocket"));
   return 0;
}

int SSLSocket::ReceiveFromTimeout(int timeout, u8_t * pBuf, int bufLen, int nFlags, CFString &peerAddress, u16_t &peerPort)
{
   throw SSLSocketException(_S("Don't call 'ReceiveFromTimeout' from a SSLSocket"));
   return 0;
}

int SSLSocket::SendTo(const u8_t * pBuf, int bufLen, int nFlags, const CFString &peerAddress, u16_t peerPort)
{
   throw SSLSocketException(_S("Don't call 'SendTo' from a SSLSocket"));
   return 0;
}

//////////////////// End of not supported methods /////////////////////

/**
 *  SSLSocket is supposed to be based on TCP only
 */
bool SSLSocket::Create(const CFString &socketAddress, u16_t socketPort)
{
   if (!Socket::Create(socketAddress, socketPort, Socket::Stream))
      throw SSLSocketException(_S("Create socket failed"));
   
   if (!PostInitSSLObj())
      throw SSLSocketException(_S("Create OpenSSL BIO failed"));

   return true;
}

/**
 *  Call by SSLContext (server context) to initialize an already accepted Socket
 *
 *  @param [in] sslObj -- SSL object created from SSL context
 *  @param [in] pPeerCertCb -- callback to peek at the certificates received. 
 *                             Can be NULL if not interested
 *  @param [in] isServer -- whether this socket is a server or a client
 */
void SSLSocket::Initialize(ptrany_t sslObj, ISSLPeerCertCallback *pPeerCertCb, bool isServer)
{
   this->sslObj = sslObj;
   this->pPeerCertCb = pPeerCertCb;
   this->sockAttr->SetAttr(Socket::SA_ISSERVER, isServer);

   if (this->pPeerCertCb) 
      cppmempool::Retain(this->pPeerCertCb);

   if (IsValidSock()) {
      if (!PostInitSSLObj())
         throw SSLSocketException(_S("Create OpenSSL BIO for accepted socket failed"));
   }
}

/**
 * Init SSLObj after the underlying socket object is initialized
 */
bool SSLSocket::PostInitSSLObj()
{
   // sslObj requires a BIO object
   BIO *bioSock = BIO_new(BIO_s_socket());
   if (!bioSock)
      return false;

   // link the BIO with our socket
   BIO_set_fd(bioSock, this->GetInternalSock(), BIO_NOCLOSE);
   // give this BIO to sslObj (reading and writing with the same BIO)
   SSL_set_bio(__GetSSL(sslObj), bioSock, bioSock); // this BIO is freed when 'SSL_free' is called)

   return true;
}

/**
 *  Do some checking after SSL negotiation succeed
 */
bool SSLSocket::PostConnectionCheck(void)
{

   if (sockAttr->IsAttrSet(Socket::SSL_CHK_PEERCERT)) {

      X509 *cert = SSL_get_peer_certificate(__GetSSL(sslObj)); 
      if (cert == NULL)   // SSL_connect can succeed without peer providing a cert
         throw SSLException(_S("No certificate from peer"));

      // 
      // Further checking of peer's domain name against the common name 
      // embedded in the cert can be performed here      
      //
      X509Certificate x509Cert(cert, true);
      if (this->pPeerCertCb)
         this->pPeerCertCb->ReadPeerCertAfterConnection(x509Cert);  // this may throw exception
   }

   if (sockAttr->IsAttrSet(Socket::SSL_CHK_VERIFYRESULT)) {

      // final checking
      long err = SSL_get_verify_result(__GetSSL(sslObj));
      if (err != X509_V_OK) {
         CFString errStr("Error verifying peer: ");
         errStr += X509_verify_cert_error_string(err);
         throw SSLException(errStr);
      }

   }

   return true;
}

/**
 *  Checking error for blocking call
 *
 *  @param [in, out] ret -- return value from a SSL I/O operation
 *
 *  @return true -- retry the previous IO operation
 *  @return false -- quit
 */
bool SSLSocket::CheckSSLBlockingIOError(int &sslCallRet)
{
   const int timeout = 30 * 1000;
   return CheckSSLNonBlockingIOError(timeout, sslCallRet);
}

/**
 *  Check error for non-blocking call
 *
 *  @return true -- retry the previous IO operation
 *  @return false -- quit
 */
bool SSLSocket::CheckSSLNonBlockingIOError(int timeout, int &sslCallRet)
{
   bool isRetryIO = false;

   int err = SSL_get_error(__GetSSL(sslObj), sslCallRet);

   switch(err) {
      case SSL_ERROR_NONE:
         break; // no error
      case SSL_ERROR_WANT_READ:
         {
//printf("%s -- SSL_ERROR_WANT_READ\n", isServer ? "server" : "client");
            int ret = ChkPendingData(timeout);
            isRetryIO = (ret > 0) ? true : false;
            if (ret == 0) sslCallRet = -1;  // timeout
            else if (ret < 0) sslCallRet = 0; // other error
         }
         break;
      case SSL_ERROR_WANT_WRITE:
         {
//printf("%s -- SSL_ERROR_WANT_WRITE\n", isServer ? "server" : "client");
            int ret = ChkReadyToWrite(timeout);
            isRetryIO = (ret > 0) ? true : false;
            if (ret == 0) sslCallRet = -1;  // timeout
            else if (ret < 0) sslCallRet = 0; // other error
         }
         break;
      case SSL_ERROR_ZERO_RETURN:
         sslCallRet = 0;  // connection may be closed by peer
         break;
      default:
         sslCallRet = 0;  // other error
         break;
   }

   return isRetryIO;
}

/**
 *  Set target hostname. Since HTTPS server can only read the "Host:" header after SSL negotiation,
 *  this hostname must be set in advance or target (virtual) server doesn't know which host you want
 *  to access (Server Name Indication).
 */
void SSLSocket::SetTargetHost(const CFString & targetHost)
{
   ObjPtrT<char> pHostName(targetHost.ToUtf8String());
   if (!SSL_set_tlsext_host_name(__GetSSL(sslObj), static_cast<char*>(pHostName)))
      throw SSLException(_S("Error setting target hostname for SSL connection"));

   // assign to call back, too, if any
   if (pPeerCertCb) {
      pPeerCertCb->SetTargetHostName(targetHost);
   }
}

/**
 *  Blocking connect
 */
bool SSLSocket::Connect(const CFString &hostAddress, u16_t hostPort)
{
   // normal socket connection first
   if (!Socket::Connect(hostAddress, hostPort))
      throw SSLSocketException(_S("Connect to server failed"));

   if (!sockAttr->IsAttrSet(Socket::SSL_SENDRECV_WITHSSL))
      return true;

   SetTargetHost(hostAddress);

   // SSL negotiation
   if (SSL_connect(__GetSSL(sslObj)) <= 0) 
      throw SSLException(_S("SSL connect failed"));

   return PostConnectionCheck();
}



/**
 *  Connect with a timeout
 */
bool SSLSocket::Connect(const CFString &hostAddress, u16_t hostPort, int timeout)
{
   if (!Socket::Connect(hostAddress, hostPort, timeout))
      throw SSLSocketException(_S("Connect to server failed"));

   if (!sockAttr->IsAttrSet(Socket::SSL_SENDRECV_WITHSSL))
      return true;

   SetTargetHost(hostAddress);

   // must set to non-blocking to detect timeout or SSL_connect may block forever
   SocketToggleBlockingState toggleBlockingState(this, false);

   int ret = 0;
   do {
      ret = SSL_connect(__GetSSL(sslObj));
      if (!CheckSSLNonBlockingIOError(timeout, ret))
         break;
   } while(true);

   if (ret <= 0)
      throw SSLException(_S("SSL connect failed"));

   return PostConnectionCheck();
}


/**
 *  @return > 0 -- number of bytes received
 *  @return 0 -- connection closed by peer
 *  @return < 0 -- error
 */
int SSLSocket::Receive(u8_t * pBuf, int bufLen, int nFlags)
{
   if (!sockAttr->IsAttrSet(Socket::SSL_SENDRECV_WITHSSL))
      return Socket::Receive(pBuf, bufLen, nFlags);

   int ret = 0;
   do {
      ret = SSL_read(__GetSSL(sslObj), pBuf, bufLen);
      if (!CheckSSLBlockingIOError(ret))
         break;
   } while(true);

   return ret;
}

/**
 *  Receive data with a timeout value
 *
 *  @return > 0 -- number of bytes received
 *  @return 0 -- error or connection closed by peer
 *  @return -1 -- timeout
 */
int SSLSocket::ReceiveTimeout(int timeout, u8_t * pBuf, int bufLen, int nFlags)
{
   if (!sockAttr->IsAttrSet(Socket::SSL_SENDRECV_WITHSSL))
      return Socket::ReceiveTimeout(timeout, pBuf, bufLen, nFlags);

   // must set to non-blocking to detect timeout or SSL_read may block forever
   SocketToggleBlockingState toggleBlockingState(this, false);

   int ret = 0;
   do {
      ret = SSL_read(__GetSSL(sslObj), pBuf, bufLen);
      if (!CheckSSLNonBlockingIOError(timeout, ret))
         break;
   } while(true);
   
   return ret;
}

/**
 *  @return > 0 -- number of bytes sent
 *  @return 0 -- connection closed by peer
 *  @return < 0 -- error
 */
int SSLSocket::Send(const u8_t * pBuf, int bufLen, int nFlags)
{
   if (!sockAttr->IsAttrSet(Socket::SSL_SENDRECV_WITHSSL))
      return Socket::Send(pBuf, bufLen, nFlags);

   // must set to non-blocking to detect timeout or SSL_write may block (if a SSL_ERROR_WANT_READ is returned)
   SocketToggleBlockingState toggleBlockingState(this, false);

   int ret = 0;
   do {
      ret = SSL_write(__GetSSL(sslObj), pBuf, bufLen);
      if (!CheckSSLBlockingIOError(ret))
         break;
   } while(true);

   return ret;
}

/**
 *  Close socket
 */
void SSLSocket::Close()
{
   if (sslObj) {
      SSL_shutdown(__GetSSL(sslObj));
      SSL_free(__GetSSL(sslObj));
      ERR_remove_state(0);
      sslObj = NULL;
   }

   if (pPeerCertCb) {
      cppmempool::Release(pPeerCertCb);
      pPeerCertCb = NULL;
   }

   Socket::Close();
}

/**
 *  Waiting for handshake to start from peer
 */
void SSLSocket::WaitHandshake(int timeout)
{
   // must set to non-blocking to detect timeout or SSL_accept may block
   SocketToggleBlockingState toggleBlockingState(this, false);

   int ret = 0;
   do {
      ret = SSL_accept(__GetSSL(sslObj));
      if (!CheckSSLNonBlockingIOError(timeout, ret))
         break;
   } while(true);

   if (ret < 0)
      throw SSLException(_S("Error occur during SSL handshaking with peer"));
   else if (ret == 0)
      throw SSLException(_S("Error: peer shutdown during SSL handshaking"));

   PostConnectionCheck();
}

PRIVATE void SSL_set_state(SSL *sslObj, int state)
{
   sslObj->state = state;
}

/**
 *  Renegotiate with peer (exchange new key to encrypt data)
 *
 *  @param [in] isWaitPeer -- true to ensure the renegotiation must complete
 *                            false to indicate peer can ignore the renegotiation
 *
 *  @param [in] timeout -- timeout value if isWaitPeer is true
 *  @param [in] needPeerAuth -- (for server socket only) 
 *                              pass true to require client to provide its cert for authentication
 *                              (useful when server want to upgrade client authentication)
 *
 *  @note If it is a server socket and 'isWaitPeer' is false (i.e. don't wait peer), client may
 *        not perform renegotiation.
 *
 *  @return true OK
 *  @return false failed
 */
bool SSLSocket::Renegotiate(bool isWaitPeer, int timeout, bool needPeerAuth)
{

   if (sockAttr->IsAttrSet(Socket::SA_ISSERVER) && needPeerAuth) {
      int orgMode = SSL_get_verify_mode(__GetSSL(sslObj));
      // passing NULL will keep the original callback
      SSL_set_verify(__GetSSL(sslObj), orgMode | SSL_VERIFY_PEER | SSL_VERIFY_FAIL_IF_NO_PEER_CERT, NULL); 
   }


   SSL_renegotiate(__GetSSL(sslObj));  // set flag first
   if (!isWaitPeer)   // don't wait peer
      return true;

   do {
      // must set to non-blocking to detect timeout or SSL_do_handshake may block
      SocketToggleBlockingState toggleBlockingState(this, false);      

      int ret = 0;
      do {
         ret = SSL_do_handshake(__GetSSL(sslObj));  // send request to peer
         if (!CheckSSLNonBlockingIOError(timeout, ret))
            break;
      } while(true);

      if (ret <= 0) 
         break;
      
      if (sockAttr->IsAttrSet(Socket::SA_ISSERVER)) {
         // must set 'ACCEPT' for a server socket once and only once before the next 'SSL_do_handshake'
         SSL_set_state(__GetSSL(sslObj), SSL_ST_ACCEPT); 
      }      
      
      ret = 0;
      do {
         ret = SSL_do_handshake(__GetSSL(sslObj));  
         if (!CheckSSLNonBlockingIOError(timeout, ret))
            break;         
      } while(true);

      if (ret <= 0)
         break;

      return true;

   } while(false);

   return false;
}

/**
 *  Currently, only one attribute is allowed to be enabled/disabled
 */
void SSLSocket::SetAttribute(sockattr_t attr, const BaseObject &val)
{
   if (attr == Socket::SSL_SENDRECV_WITHSSL) {
      // enable/disable SSL overriden methods
      if (!cppflib::InstanceOf<Boolean>(&val)) 
         throw SSLException(_S("SetAttribute: value is not Boolean"));
      this->sockAttr->SetAttr(attr, static_cast<const Boolean*>(&val)->Value());
   }
   else {
      Socket::SetAttribute(attr, val);  // fall back to base class
   }
}

/**
 *  Get peer certificate after negotiation
 *
 *  @return NULL if not successful
 */
X509Certificate *SSLSocket::GetPeerCertificate()
{
   if (!sslObj)
      return NULL;

   X509 *cert = SSL_get_peer_certificate(__GetSSL(sslObj)); 
   if (cert == NULL)
      return NULL;
      
   return new X509Certificate(cert, true);
}

} // end of namespace ssl

} // end of namespace net

} // end of namespace cppflib
